from __future__ import annotations

import pathlib, os
from typing import Mapping
from google.protobuf.compiler import plugin_pb2 as plugin
from google.protobuf.descriptor_pool import DescriptorPool
from .options import parse_options
from .templates import TemplateRenderer
from .descriptor_ext import ProtoFile
from .generator import generate_single_file, output_filename
from .utils import merge_dts


def generate_code(request: plugin.CodeGeneratorRequest) -> plugin.CodeGeneratorResponse:
    opts = parse_options(request.parameter)
    response = plugin.CodeGeneratorResponse(supported_features=plugin.CodeGeneratorResponse.FEATURE_SUPPORTS_EDITIONS)
    if opts.help:
        return response

    # 预解析所有文件
    pool = DescriptorPool()
    file_map: Mapping[str, ProtoFile] = {fd.name: ProtoFile(fd, opts) for fd in request.proto_file}
    for fd in request.proto_file:
        pool.Add(fd)
    renderer = TemplateRenderer(pool, opts)
    all_in_one_init = []
    all_dts_path: list[pathlib.Path] = []
    symbol_map: dict[ProtoFile, tuple[list[str], list[str]]] = {}
    for file_to_gen in request.file_to_generate:
        proto_file = file_map[file_to_gen]
        header_content, cc_content, dts_content = generate_single_file(proto_file, file_map, renderer)
        header_name, cc_name, dts_name = output_filename(file_to_gen)
        response.file.add(name=header_name, content=header_content)
        response.file.add(name=cc_name, content=cc_content)
        all_in_one_init.append((header_name, proto_file))

        dts_path = pathlib.Path(opts.dts_out) / dts_name
        os.makedirs(dts_path.parent, exist_ok=True)
        dts_path.write_text(dts_content)
        all_dts_path.append(dts_path)
        symbol_map[proto_file] = (renderer.enums, renderer.classes)

    response.file.add(name="napi_init.h", content=renderer.all_in_one_init(all_in_one_init))

    napi_utils_h = (pathlib.Path(__file__).parent.parent / "cpp/napi_utils.h").read_text()
    if opts.enable_json:
        napi_utils_h = "#define NAPI_UTILS_ENABLE_JSON\n\n" + napi_utils_h
    response.file.add(name="napi_utils.h", content=napi_utils_h)
    napi_utils_cpp = (pathlib.Path(__file__).parent.parent / "cpp/napi_utils.cpp").read_text()
    response.file.add(name="napi_utils.cpp", content=napi_utils_cpp)
    export_other_dts = [
        f"export * from './{dts_path.relative_to(opts.dts_out).as_posix()[:-5]}';" for dts_path in all_dts_path
    ] if not opts.enable_namespace else [merge_dts(opts.dts_out)]
    global_functions = []
    if opts.enable_arena_switch:
        global_functions.append("export const setArenaEnabled: (enabled: boolean) => void;")
    if opts.enable_json:
        global_functions.append("export const toJsonString: (obj: object) => string;")
    all_in_one_dts_path = pathlib.Path(opts.dts_out) / "all_in_one.d.ts"
    all_in_one_dts_content = "\n".join(
        [
            "// Generated by protoc-gen-napi. DO NOT EDIT!",
            *global_functions,
            *export_other_dts,
        ]
    )

    all_in_one_dts_path.write_text(all_in_one_dts_content)
    if opts.index_ts and not opts.enable_namespace:
        content = "\n".join(renderer.index_ts(symbol_map))
        index_ts_path = pathlib.Path(opts.index_ts)
        os.makedirs(index_ts_path.parent, exist_ok=True)
        index_ts_path.write_text(content)

    return response
